#!/usr/bin/env python3

import asyncio
import gc
import logging
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, List, Optional

# 尝试导入不同的依赖
try:
    from modelscope.pipelines import pipeline
    from modelscope.utils.config import Config
    MODELSCOPE_AVAILABLE = True
    FUNASR_AVAILABLE = False
except ImportError:
    try:
        from funasr import AutoModel
        FUNASR_AVAILABLE = True
        MODELSCOPE_AVAILABLE = False
        logging.warning("ModelScope未安装，但FunASR可用。注意：FunASR不支持NLP模型")
    except ImportError:
        logging.warning("ModelScope和FunASR都未安装，请先安装其中之一")
        MODELSCOPE_AVAILABLE = False
        FUNASR_AVAILABLE = False

from ...config.funasr_config import ModelConfig
from ..schemas import ActionClassification

logger = logging.getLogger(__name__)

class StructBertManager:
    """
    StructBERT模型管理器（性能优化版本）

    负责管理StructBERT模型的加载、推理和资源管理
    用于会议文本的行动分类和关键片段识别
    包含性能优化和内存管理功能
    """

    def __init__(self):
        self.model = None
        self.model_loaded = False
        self.model_name = ModelConfig.ACTION_MODEL
        self.model_revision = ModelConfig.ACTION_MODEL_REVISION
        self.load_time = None
        self.last_inference_time = None

        # 线程池优化
        self.executor = ThreadPoolExecutor(
            max_workers=3,
            thread_name_prefix="structbert"
        )

        # 性能统计
        self.stats = {
            'total_inferences': 0,
            'total_inference_time': 0,
            'average_inference_time': 0,
            'load_attempts': 0,
            'successful_loads': 0,
            'errors': 0
        }

        # 分类标签映射（优化后的映射）
        self.label_mapping = {
            0: "讨论",
            1: "决策",
            2: "任务",
            3: "问题",
            4: "总结",
            5: "其他"
        }

        # 反向映射，用于快速查找
        self.reverse_label_mapping = {v: k for k, v in self.label_mapping.items()}

        # 模型配置（性能优化）
        self.config = {
            'model': self.model_name,
            'model_revision': self.model_revision,
            'device': 'cpu',  # 可以根据配置调整为'gpu'
            'batch_size': 8,  # 批处理大小
            'max_seq_length': 512,  # 最大序列长度
            'num_workers': 2,  # 数据加载器工作线程数
        }

        # 文本处理配置
        self.text_config = {
            'max_segment_length': 200,  # 最大片段长度
            'overlap_length': 20,       # 重叠长度
            'min_segment_length': 10,   # 最小片段长度
        }

    async def load_model(self) -> bool:
        """
        异步加载模型（使用funasr的AutoModel）

        Returns:
            bool: 加载是否成功
        """
        if self.model_loaded and self.model is not None:
            logger.info("StructBERT模型已加载")
            return True

        # 优先使用ModelScope（因为这是NLP模型）
        if MODELSCOPE_AVAILABLE:
            return await self._load_with_modelscope()
        else:
            logger.error("ModelScope未安装，无法加载NLP模型")
            logger.info("提示：请安装ModelScope：pip install modelscope")
            return False

    async def _load_with_modelscope(self) -> bool:
        """使用ModelScope加载模型"""

        self.stats['load_attempts'] += 1
        start_time = time.time()

        try:
            logger.info(f"开始加载StructBERT模型: {self.model_name}")

            # 使用新的模型加载器
            from .modelscope_loader import get_modelscope_loader
            loader = get_modelscope_loader()

            # 在线程池中加载模型以避免阻塞
            loop = asyncio.get_event_loop()
            self.model = await loop.run_in_executor(
                self.executor,
                self._load_model_sync_modelscope,
                loader
            )

            if self.model is not None:
                self.model_loaded = True
                self.load_time = time.time() - start_time
                self.stats['successful_loads'] += 1
                logger.info(f"StructBERT模型加载成功，耗时: {self.load_time:.2f}秒")
                return True
            else:
                logger.error("StructBERT模型加载失败")
                return False

        except Exception as e:
            self.stats['errors'] += 1
            logger.error(f"StructBERT模型加载异常: {e}")
            return False

    def _load_model_sync_modelscope(self, loader):
        """使用ModelScope同步加载模型"""
        try:
            # 使用新的加载器从本地加载模型
            model = loader.load_model_with_fallback(
                model_type="structbert",
                task='text-classification',
                device=self.config['device']
            )
            return model
        except Exception as e:
            logger.error(f"ModelScope同步加载模型失败: {e}")
            return None

    async def classify_text(
        self,
        text: str,
        segment_length: Optional[int] = None
    ) -> Optional[ActionClassification]:
        """
        对文本进行分类（性能优化版本）

        Args:
            text: 输入文本
            segment_length: 分段长度

        Returns:
            Optional[ActionClassification]: 分类结果
        """
        if not self.model_loaded or self.model is None:
            logger.warning("模型未加载，尝试加载模型")
            if not await self.load_model():
                logger.error("模型加载失败，无法进行分类")
                return None

        start_time = time.time()

        try:
            # 文本预处理和分段
            segments = self._split_text_optimized(text, segment_length)

            if not segments:
                logger.warning("文本分段为空")
                return None

            # 批量推理（如果有多个段落）
            if len(segments) > 1:
                classification_result = await self._batch_classify_segments(segments)
            else:
                classification_result = await self._classify_single_segment(segments[0])

            # 更新统计信息
            inference_time = time.time() - start_time
            self.last_inference_time = inference_time
            self.stats['total_inferences'] += 1
            self.stats['total_inference_time'] += inference_time
            self.stats['average_inference_time'] = (
                self.stats['total_inference_time'] / self.stats['total_inferences']
            )

            logger.info(f"文本分类完成，耗时: {inference_time:.3f}秒")
            return classification_result

        except Exception as e:
            self.stats['errors'] += 1
            logger.error(f"文本分类失败: {e}")
            return None

    def _split_text_optimized(self, text: str, segment_length: Optional[int] = None) -> List[str]:
        """
        优化的文本分割方法

        Args:
            text: 输入文本
            segment_length: 分段长度

        Returns:
            List[str]: 文本段落列表
        """
        if not text or not text.strip():
            return []

        text = text.strip()
        max_length = segment_length or self.text_config['max_segment_length']
        min_length = self.text_config['min_segment_length']
        overlap = self.text_config['overlap_length']

        # 如果文本长度小于最大长度，直接返回
        if len(text) <= max_length:
            return [text] if len(text) >= min_length else []

        segments = []
        start = 0

        while start < len(text):
            end = min(start + max_length, len(text))

            # 尝试在句号、感叹号或问号处切分
            if end < len(text):
                for punct in ['。', '！', '？', '.', '!', '?']:
                    punct_pos = text.rfind(punct, start, end)
                    if punct_pos > start + min_length:
                        end = punct_pos + 1
                        break

            segment = text[start:end].strip()
            if len(segment) >= min_length:
                segments.append(segment)

            # 计算下一个起始位置（考虑重叠）
            if end >= len(text):
                break
            start = max(start + 1, end - overlap)

        return segments

    async def _classify_single_segment(self, segment: str) -> Optional[ActionClassification]:
        """
        分类单个文本段落

        Args:
            segment: 文本段落

        Returns:
            Optional[ActionClassification]: 分类结果
        """
        try:
            # 在线程池中执行推理，使用funasr的generate方法
            loop = asyncio.get_event_loop()
            result = await loop.run_in_executor(
                self.executor,
                self.model,
                segment
            )

            if result and isinstance(result, dict):
                label = result.get('label', 'LABEL_0')
                score = result.get('score', 0.0)

                # 解析标签
                label_text = self._parse_label(label)

                return ActionClassification(
                    label=label_text,
                    confidence=float(score),
                    text_segment=segment[:100] + "..." if len(segment) > 100 else segment
                )
            else:
                logger.warning(f"推理结果格式异常: {result}")
                return None

        except Exception as e:
            logger.error(f"单段落分类失败: {e}")
            return None



    async def _batch_classify_segments(self, segments: List[str]) -> Optional[ActionClassification]:
        """
        批量分类多个文本段落并合并结果

        Args:
            segments: 文本段落列表

        Returns:
            Optional[ActionClassification]: 合并后的分类结果
        """
        try:
            # 并发处理多个段落
            tasks = [
                self._classify_single_segment(segment)
                for segment in segments
            ]

            results = await asyncio.gather(*tasks, return_exceptions=True)

            # 过滤有效结果
            valid_results = [
                r for r in results
                if isinstance(r, ActionClassification) and r is not None
            ]

            if not valid_results:
                logger.warning("没有有效的分类结果")
                return None

            # 合并结果（投票机制）
            return self._merge_classifications(valid_results, segments)

        except Exception as e:
            logger.error(f"批量分类失败: {e}")
            return None

    def _merge_classifications(
        self,
        classifications: List[ActionClassification],
        original_segments: List[str]
    ) -> ActionClassification:
        """
        合并多个分类结果

        Args:
            classifications: 分类结果列表
            original_segments: 原始文本段落

        Returns:
            ActionClassification: 合并后的分类结果
        """
        # 统计各标签的票数和置信度
        label_votes = {}
        label_confidences = {}

        for cls in classifications:
            label = cls.label
            confidence = cls.confidence

            if label not in label_votes:
                label_votes[label] = 0
                label_confidences[label] = []

            label_votes[label] += 1
            label_confidences[label].append(confidence)

        # 找到得票最多的标签
        winning_label = max(label_votes.keys(), key=lambda x: label_votes[x])

        # 计算平均置信度
        avg_confidence = sum(label_confidences[winning_label]) / len(label_confidences[winning_label])

        # 创建合并的文本段落
        merged_text = "...".join([seg[:50] for seg in original_segments[:3]])
        if len(merged_text) > 200:
            merged_text = merged_text[:200] + "..."

        return ActionClassification(
            label=winning_label,
            confidence=avg_confidence,
            text_segment=merged_text
        )

    def _parse_label(self, raw_label: str) -> str:
        """
        解析模型输出的标签

        Args:
            raw_label: 原始标签

        Returns:
            str: 解析后的标签文本
        """
        try:
            # 处理 LABEL_0, LABEL_1 格式
            if raw_label.startswith('LABEL_'):
                label_id = int(raw_label.split('_')[1])
                return self.label_mapping.get(label_id, "其他")

            # 直接返回标签文本
            return raw_label if raw_label in self.reverse_label_mapping else "其他"

        except Exception as e:
            logger.warning(f"标签解析失败: {e}, 原始标签: {raw_label}")
            return "其他"

    async def get_model_status(self) -> Dict[str, Any]:
        """
        获取模型状态（增强版本）

        Returns:
            Dict[str, Any]: 模型状态信息
        """
        try:
            return {
                'model_name': self.model_name,
                'model_revision': self.model_revision,
                'loaded': self.model_loaded,
                'load_time': self.load_time,
                'last_inference_time': self.last_inference_time,
                'statistics': self.stats.copy(),
                'config': self.config.copy(),
                'available_labels': list(self.label_mapping.values()),
                'funasr_available': FUNASR_AVAILABLE
            }
        except Exception as e:
            logger.error(f"获取模型状态失败: {e}")
            return {
                'loaded': False,
                'error': str(e)
            }

    async def warm_up(self, sample_texts: Optional[List[str]] = None) -> bool:
        """
        模型预热

        Args:
            sample_texts: 样本文本列表

        Returns:
            bool: 预热是否成功
        """
        if not self.model_loaded:
            logger.warning("模型未加载，无法预热")
            return False

        try:
            # 使用默认样本文本进行预热
            if not sample_texts:
                sample_texts = [
                    "这是一个测试文本，用于模型预热。",
                    "我们需要讨论一下项目的进展情况。",
                    "请大家按时完成分配的任务。"
                ]

            logger.info("开始模型预热...")
            start_time = time.time()

            # 并发执行预热推理
            tasks = [
                self._classify_single_segment(text)
                for text in sample_texts
            ]

            results = await asyncio.gather(*tasks, return_exceptions=True)

            successful_warm_ups = sum(
                1 for r in results
                if isinstance(r, ActionClassification)
            )

            warm_up_time = time.time() - start_time
            logger.info(f"模型预热完成，成功{successful_warm_ups}/{len(sample_texts)}次，耗时: {warm_up_time:.2f}秒")

            return successful_warm_ups > 0

        except Exception as e:
            logger.error(f"模型预热失败: {e}")
            return False

    async def cleanup(self):
        """
        清理资源（增强版本）
        """
        logger.info("开始清理StructBERT资源...")

        try:
            # 清理模型
            if hasattr(self, 'model') and self.model is not None:
                del self.model
                self.model = None

            self.model_loaded = False

            # 关闭线程池
            if hasattr(self, 'executor') and self.executor:
                self.executor.shutdown(wait=True)

            # 执行垃圾回收
            collected = gc.collect()
            logger.info(f"StructBERT资源清理完成，回收了{collected}个对象")

        except Exception as e:
            logger.error(f"清理StructBERT资源时出错: {e}")

    def reset_stats(self):
        """重置统计信息"""
        self.stats = {
            'total_inferences': 0,
            'total_inference_time': 0,
            'average_inference_time': 0,
            'load_attempts': 0,
            'successful_loads': 0,
            'errors': 0
        }
        logger.info("StructBERT统计信息已重置")

    def update_config(self, new_config: Dict[str, Any]):
        """
        更新配置

        Args:
            new_config: 新配置
        """
        old_config = self.config.copy()
        self.config.update(new_config)
        logger.info(f"StructBERT配置已更新: {new_config}")

        # 如果设备配置改变，需要重新加载模型
        if old_config.get('device') != self.config.get('device'):
            logger.info("设备配置改变，需要重新加载模型")
            self.model_loaded = False
            self.model = None
